from pydantic import BaseModel, Field
from langchain.agents import tool

from utils.utils import extract_data, jsonOutputParser
import json
from config import web_json
from llm.llm import llm

from langchain.prompts import (
    SystemMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)


def overview_style(str):
    sys_tmp = """\
        # 角色设定
        你是一名专业的Web前端工程师，擅长数据可视化设计和UI样式配置。

        # 任务要求
        请为提供的JSON数据中的数组的每个项添加适当的样式配置，包括：
        1. `icon` - 根据内容选择匹配的图标（使用Font Awesome类名）
        2. `background` - 创建美观的渐变色背景（使用CSS linear-gradient语法）
        3. `color` - 选择与背景对比度良好的文字颜色（十六进制格式）

        # 设计原则
        1. 保持视觉一致性，相同type的数据使用相似的配色方案
        2. 图标选择应语义化匹配字段含义
        3. 文字颜色需确保在背景上清晰可读
        4. 渐变背景应美观协调，避免使用冲突色

        # 输出要求
        - 必须返回严格规范的JSON格式
        - 每个数据项必须包含完整的style对象
        - 图标字段留空时使用空字符串""

        # 输入示例
        ```json
        [
            {{
                "label": "名称",
                "value": "服务器集群A"
            }}
        ]

        返回JSON结果
        ```json
        [
            {{
                "label": "名称",
                "value": "服务器集群A",
                "style": {{
                    "color": "#ffffff",
                    "background": "linear-gradient(135deg, #3a4a9f 0%, #1a2a6c 100%)",
                    "icon": "fa-server"
                }}
            }}
        ]
        ```
    """

    human_tmp = """\
        待处理数据
        {data}
    """

    sys_msg = SystemMessagePromptTemplate.from_template(sys_tmp)
    human_msg = HumanMessagePromptTemplate.from_template(human_tmp)
    chat_msg = ChatPromptTemplate.from_messages([sys_msg, human_msg])

    res = ""
    try:
        chain = (chat_msg | llm | jsonOutputParser()).with_retry()
        res = chain.invoke(
            {
                "data": str,
            },
        )
        print("处理结果", res)

    except Exception as e:
        print("异常", e)
        return "数据分析失败，请重试"

    return res


class overview_input(BaseModel):
    title: str = Field(..., description="生成概览组件的标题")
    desc: str = Field(
        ...,
        description="适合概览展示的文本段落",
    )


overview_desc = """\
提取适合生成概览信息的文本段落

要求：
1. 文本应包含系统或数据的整体情况描述
2. 可以包含关键指标或统计数字
3. 适合快速了解全局情况的内容
"""


@tool(description=overview_desc, args_schema=overview_input)
def overview(*args, **kwargs) -> str:
    title = ""
    if "title" in kwargs:
        title = kwargs["title"]
    else:
        return "输入参数错误，根据工具参数要求输入title参数"

    desc = ""
    if "desc" in kwargs:
        desc = kwargs["desc"]
    else:
        return "输入参数错误，根据工具参数要求输入desc参数"

    data = extract_data(desc)
    data = overview_style(data)

    res = {
        "title": title,
        "type": "overview",
        "data": json.loads(data),
    }

    web_json.append(res)

    # return json.dumps(res, ensure_ascii=False)
    return "概览生成成功"
